import os
import math
import pylab as plt
import numpy as np
from params import *

datasets = ['bkite','fsquare','livejournal', 'twitter']
lbls = ['Brightkite', 'Foursquare', 'LiveJournal', 'Twitter']

for dataset in datasets:
  clustfile  = os.path.join(WORKDIR, 'gsn', 'results', dataset, '%s_local_clustering.txt'%dataset)
  locfile = os.path.join(WORKDIR, 'gsn', 'results',dataset, '%s_node_locality.txt'%dataset)

  if dataset in ['bkite', 'fsquare']:
      directed = 0
  else:
      directed = 1

  loc_clust = {}
  for line in open(clustfile):
      user, clust = line.split(' ')
      loc_clust[int(user)] = float(clust)

  in_deg = {}
  out_deg = {}
  for line in open(locfile):
      user = int(line.split(' ')[0])
      outdeg,dout,indeg,din = map(float,line.split(' ')[1:])
      in_deg[user] = indeg
      out_deg[user] = outdeg

  print "read %d users"%len(in_deg)
  data = [(in_deg[u], loc_clust[u]) for u in loc_clust]
  data2 = [(out_deg[u], loc_clust[u]) for u in loc_clust]
  m = math.ceil(math.log10(max(in_deg[u] for u in in_deg)))
  bins = 10**np.linspace(0,m,1000)
  print m
  plt.figure()
  y_bins = {}
  for d,l in data:
      i=0
      while d > bins[i]:
          i+=1

      d = bins[i]
      if not d in y_bins:
          y_bins[d] = (0.0, 0)
      s, c = y_bins[d]
      s += l
      c += 1
      y_bins[d] = (s,c)

  x = sorted(y_bins.keys())
  y = [y_bins[d][0]/y_bins[d][1]  for d in x]

  y_bins = {}
  for d,l in data2:
      i=0
      while d > bins[i]:
          i+=1

      d = bins[i]
      if not d in y_bins:
          y_bins[d] = (0.0, 0)
      s, c = y_bins[d]
      s += l
      c += 1
      y_bins[d] = (s,c)

  x2 = sorted(y_bins.keys())
  y2 = [y_bins[d][0]/y_bins[d][1]  for d in x2]

  ylab ='Geographic clustering'
  xlab = 'Degree'
  name = '%s_geoclustering_degree_correlation.pdf'%dataset
  plt.figure()
  plt.clf()
  plt.axes(FIG_AXES2)
  plt.semilogx(x,y,'kx')
  plt.grid(True)
  plt.xlabel(xlab)
  plt.ylabel(ylab)
  if directed:
      plt.semilogx(x2,y2,'ko', mfc='None')
      plt.legend(['In', 'Out'], loc='upper right',numpoints=1)
  tick,lbl = plt.yticks()
  tick = [0.1*i for i in range(5)]
  lbl = map(lambda x: '%.1f'%x,tick)
  plt.yticks(tick,lbl)
  plt.axis([1,10**m,0,0.4])
  #plt.axis([0.001,10**4,0,1])
  plt.savefig(name)
  plt.close()

  continue

  def plot_correlation(data, xlab, ylab, name,l1,l2):
      bins = 10**np.linspace(l1,l2,1000)
      plt.figure()
      y_bins = {}
      for d,l in data:
          i=0
          while d > bins[i]:
              i+=1

          d = bins[i]
          if not d in y_bins:
              y_bins[d] = (0.0, 0)
          s, c = y_bins[d]
          s += l
          c += 1
          y_bins[d] = (s,c)

      x = sorted(y_bins.keys())
      y = [y_bins[d][0]/y_bins[d][1]  for d in x]
      plt.loglog(x,y,'k.')
      plt.grid(True)
      plt.xlabel(xlab)
      plt.ylabel(ylab)
      #plt.axis([1,10**6,0,1])
      #plt.axis([0.001,10**4,0,1])
      plt.savefig(name)
      plt.close()

  data = [(in_deg[u], loc_clust[u]) for u in loc_clust]
  if directed:
      xlab = 'Average node in-degree'
  else:
      xlab = 'Average node degree'
  ylab = 'Geographic clustering'
  name = '%s_geoclustindegree_correlation.pdf'%dataset
  plot_correlation(data,xlab,ylab,name,0,6)

  if directed:
      data = [(out_deg[u], loc_clust[u]) for u in loc_clust]
      xlab = 'Average node out-degree'
      ylab = 'Geographic clustering'
      name = '%s_geoclustoutdegree_correlation.pdf'%dataset
      plot_correlation(data,xlab,ylab,name,0,6)

      data = [(ratio[u], loc_clust[u]) for u in loc_clust if u in ratio]
      xlab = 'In-degree/Out-degree ratio'
      ylab = 'Geographic clustering'
      name = '%s_geoclustratio_correlation.pdf'%dataset
      plot_correlation(data,xlab,ylab,name,-3,5)
  sys.exit()





  data = [(ratio[u], out_loc[u]) for u in ratio]
  xlab = 'In-degree/Out-degree ratio'
  ylab = 'Average node out-locality'
  name = '%s_ratiooutloc_correlation.pdf'%dataset
  plot_correlation(data,xlab,ylab,name)

  sys.exit()

  data = [(in_deg[u], in_loc[u]) for u in in_loc]
  if directed:
      xlab = 'In-degree'
      ylab ='Average node in-locality'
  else:
      xlab = 'Degree'
      ylab ='Average node locality'
  name = '%s_ininlocdegree_correlation.pdf'%dataset
  plot_correlation(data,xlab,ylab,name)

  data = [(out_deg[u], out_loc[u]) for u in in_loc]
  if directed:
      xlab = 'Out-degree'
      ylab ='Average node out-locality'
  else:
      xlab = 'Degree'
      ylab ='Average node locality'
  name = '%s_outoutlocdegree_correlation.pdf'%dataset
  plot_correlation(data,xlab,ylab,name)

  data = [(out_deg[u], in_loc[u]) for u in in_loc]
  if directed:
      xlab = 'Out-degree'
      ylab ='Average node in-locality'
  else:
      xlab = 'Degree'
      ylab ='Average node locality'
  name = '%s_outinlocdegree_correlation.pdf'%dataset
  plot_correlation(data,xlab,ylab,name)

  data = [(in_deg[u], out_loc[u]) for u in in_loc]
  if directed:
      xlab = 'In-degree'
      ylab ='Average node out-locality'
  else:
      xlab = 'Degree'
      ylab ='Average node locality'
  name = '%s_inoutlocdegree_correlation.pdf'%dataset
  plot_correlation(data,xlab,ylab,name)
